import os
import cv2
import random
import numpy as np

image_extensions = ('.jpg', '.jpeg', '.png', '.bmp', '.tiff')

def RandomSandDust(img):
    beta = random.uniform(0.5, 1.5)
    dx = random.uniform(0.1, 1.5)
    tx1 = np.exp(-beta * dx)
    sand_dust_option = [["lightgoldenrod", 0.5098,0.86667,0.93333],
                    ["goldenrod",0.12549,0.64706,0.8549],
                    ["darkgoldenrod",0.04314,0.52549,0.72157],
                    ["peru",0.24706,0.52157,0.80392],
                    ["burlywood",0.52941,0.72157,0.87059],
                    ["wheat",0.70196,0.87059,0.96078],
                    ["sandybrown",0.37647,0.64314,0.95686],
                    ["tan",0.54902,0.70588,0.82353],
                    ["chocolate",0.11765,0.41176,0.82353],
                    ["orange",0,0.64706,1],
                    ["darkorange",0,0.54902,1],
                    ["peachpuff2",0.67843,0.79608,0.93333],
                    ["peachpuff4",0.39608,0.46667,0.5451],
                    ["navajowhite2",0.63137,0.81176,0.93333],
                    ["navajowhite3",0.5451,0.70196,0.80392]
    ]
    sand_dust_type = random.choice(sand_dust_option) #random choice one sand dust type
    # sand_dust_type = ["lightgoldenrod", 0.5098,0.86667,0.93333]
    A = sand_dust_type[1:4]
    m = img.shape[0]
    n = img.shape[1]
    rep_atmosphere = np.tile(np.reshape(A, [1, 1, 3]), [m, n, 1])
    tx1 = np.tile(tx1, [m, n, 1])
    max_transmission = np.tile(tx1, [1, 1, 3])
    img_fog = (img * max_transmission/255 + rep_atmosphere * (1 - max_transmission))*255
    return img_fog.astype(np.uint8)

def SandDust(img, sand_dust_type):
    beta = 1.0
    dx = 0.75
    tx1 = np.exp(-beta * dx)

    # sand_dust_type = ["lightgoldenrod", 0.5098,0.86667,0.93333]
    A = sand_dust_type[1:4]
    m = img.shape[0]
    n = img.shape[1]
    rep_atmosphere = np.tile(np.reshape(A, [1, 1, 3]), [m, n, 1])
    tx1 = np.tile(tx1, [m, n, 1])
    max_transmission = np.tile(tx1, [1, 1, 3])
    img_fog = (img * max_transmission/255 + rep_atmosphere * (1 - max_transmission))*255

    return img_fog.astype(np.uint8)


def inference_randomsanddust(input, batch_inference_option, input_dir, output_dir):
    # input = np.array(input)
    # import pdb;pdb.set_trace()

    if batch_inference_option:
        print(f'Batch inferencing...')
        filelist = os.listdir(input_dir)
        transformed_imagepaths = list()
        for file in filelist:
            if file.lower().endswith(image_extensions):
                filepath = os.path.join(input_dir, file)
                outputpath = os.path.join(output_dir, file)
                original_image = cv2.imread(filepath)
                transformed_image = RandomSandDust(original_image)
                cv2.imwrite(outputpath, transformed_image)
                transformed_imagepaths.append(outputpath)
        print(f'Total-{len(filelist)} sand-dusted!')
        return transformed_imagepaths

    input = cv2.cvtColor(input, cv2.COLOR_BGR2RGB)
    transformed_image = RandomSandDust(input)
    transformed_image = cv2.cvtColor(transformed_image, cv2.COLOR_RGB2BGR)

    return [transformed_image]


def inference_sanddust(input, batch_inference_option=False, input_dir=None, output_dir=None):
    sand_dust_option = [["lightgoldenrod", 0.5098,0.86667,0.93333],
                    ["goldenrod",0.12549,0.64706,0.8549],
                    ["darkgoldenrod",0.04314,0.52549,0.72157],
                    ["peru",0.24706,0.52157,0.80392],
                    ["burlywood",0.52941,0.72157,0.87059],
                    ["wheat",0.70196,0.87059,0.96078],
                    ["sandybrown",0.37647,0.64314,0.95686],
                    ["tan",0.54902,0.70588,0.82353],
                    ["chocolate",0.11765,0.41176,0.82353],
                    ["orange",0,0.64706,1],
                    ["darkorange",0,0.54902,1],
                    ["peachpuff2",0.67843,0.79608,0.93333],
                    ["peachpuff4",0.39608,0.46667,0.5451],
                    ["navajowhite2",0.63137,0.81176,0.93333],
                    ["navajowhite3",0.5451,0.70196,0.80392],
    ]

    if batch_inference_option:
        print(f'Batch inferencing...')
        filelist = os.listdir(input_dir)
        transformed_imagepaths = list()
        for file in filelist:
            if file.lower().endswith(image_extensions):
                filepath = os.path.join(input_dir, file)
                for sand_dust_type in sand_dust_option:
                    outputpath = os.path.join(output_dir, sand_dust_type[0] + "_" + file)
                    original_image = cv2.imread(filepath)
                    transformed_image = SandDust(original_image, sand_dust_type)
                    cv2.imwrite(outputpath, transformed_image)
                    transformed_imagepaths.append(outputpath)
        print(f'Total-{len(filelist)} sand-dusted!')
        return transformed_imagepaths
    
    input = cv2.cvtColor(input, cv2.COLOR_BGR2RGB)
    res_images = []
    for sand_dust_type in sand_dust_option:
        transformed_image = SandDust(input, sand_dust_type)
        transformed_image = cv2.cvtColor(transformed_image, cv2.COLOR_RGB2BGR)
        res_images.append(transformed_image)

    return res_images


from algorithms.inferencer_base import InferencerBase
class SandDustInferencer(InferencerBase):
    def __init__(self, model_name='SandDust'):
        super(SandDustInferencer, self).__init__(model_name=model_name)
        self.model = SandDust
        self.sand_dust_option = [
            ["lightgoldenrod", 0.5098,0.86667,0.93333],
            ["goldenrod",0.12549,0.64706,0.8549],
            ["darkgoldenrod",0.04314,0.52549,0.72157],
            ["peru",0.24706,0.52157,0.80392],
            ["burlywood",0.52941,0.72157,0.87059],
            ["wheat",0.70196,0.87059,0.96078],
            ["sandybrown",0.37647,0.64314,0.95686],
            ["tan",0.54902,0.70588,0.82353],
            ["chocolate",0.11765,0.41176,0.82353],
            ["orange",0,0.64706,1],
            ["darkorange",0,0.54902,1],
            ["peachpuff2",0.67843,0.79608,0.93333],
            ["peachpuff4",0.39608,0.46667,0.5451],
            ["navajowhite2",0.63137,0.81176,0.93333],
            ["navajowhite3",0.5451,0.70196,0.80392],
        ]

    def inference(self, input_data):
        input = cv2.cvtColor(input_data, cv2.COLOR_BGR2RGB)
        res_images = []
        for sand_dust_type in self.sand_dust_option:
            transformed_image = SandDust(input, sand_dust_type)
            transformed_image = cv2.cvtColor(transformed_image, cv2.COLOR_RGB2BGR)
            res_images.append(transformed_image)

        return res_images


